# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import wrapt

from experimental.prophet.graph import graph_def_pb2


class Shape(wrapt.ObjectProxy):
  @staticmethod
  def from_proto(shape_msg):
    if shape_msg.unknown_rank:
      return Shape()
    return Shape(shape_msg.dim)

  @staticmethod
  def Scalar():
    return Shape(tuple())

  @staticmethod
  def Any():
    return Shape()

  def __init__(self, *args):
    if not args:
      super().__init__(None)

    elif len(args) == 1 and isinstance(args[0], Shape):
      super().__init__(args[0].__wrapped__)

    else:
      super().__init__(tuple(*args))
      for d in self:
        if not isinstance(d, int) or d < -1:
          raise ValueError('Dimension size must be either positive, 0 or -1')

  @property
  def tuple(self):
    return self.__wrapped__

  @property
  def rank(self):
    if self.unknown_rank:
      return None
    else:
      return len(self.__wrapped__)

  @property
  def unknown_rank(self):
    return self.__wrapped__ is None

  @property
  def is_scalar(self):
    return self.rank == 0

  def is_fully_defined(self):
    if self.unknown_rank:
      return False
    for d in self:
      if d == -1:
        return False
    return True

  def __str__(self):
    if self.unknown_rank:
      return '?'
    else:
      if len(self) == 1:
        d = self[0]
        return '(%s,)' % (str(d) if d != -1 else '?')
      else:
        return '(%s)' % ', '.join(((str(d) if d != -1 else '?') for d in self))

  def __repr__(self):
    return 'Shape(%s)' % str(self)

  def to_proto(self):
    msg = graph_def_pb2.ShapeProto()
    if self.unknown_rank:
      msg.unknown_rank = True
      return msg
    msg.dim.extend(self.__wrapped__)
    return msg
